import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

# ---------------- 全局参数配置 ----------------
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "Times New Roman"]
plt.rcParams["axes.unicode_minus"] = False
plt.rcParams.update({
    "font.size": 12,
    "axes.labelsize": 14,
    "axes.titlesize": 16,
    "legend.fontsize": 11,
    "grid.linestyle": "--",
    "grid.alpha": 0.7
})

# 处理LaTeX兼容
try:
    plt.rcParams["text.usetex"] = True
    plt.rcParams["text.latex.preamble"] = [
        r"\usepackage{amsmath}",
        r"\usepackage{amsfonts}",
        r"\usepackage{amsbsy}"
    ]
except Exception:
    plt.rcParams["text.usetex"] = False


def plot_smoke_bomb_trajectory():
    """绘制烟幕弹三阶段运动轨迹图"""

    # ---------------- 参数定义 ----------------
    t_t, t_q = 10, 15          # 投放与起爆时间
    t_ground = 35              # 理论触地时间（安全冗余）
    total_duration = t_ground + 5

    v_u = 100                  # 无人机水平速度 (m/s)
    theta = 30                 # 无人机飞行方位角 (度)
    theta_rad = np.radians(theta)

    g = 9.8                    # 重力加速度 (m/s^2)
    v_sink = 5                 # 起爆后下沉速度 (m/s)
    ground_z = 0               # 地面高度

    # ---------------- 阶段1：投放前 ----------------
    t1 = np.linspace(0, t_t, 100)
    x1 = 17800 + v_u * t1 * np.cos(theta_rad)
    y1 = 0 + v_u * t1 * np.sin(theta_rad)
    z1 = 1800 + 0.5 * t1   # 假设无人机爬升

    # ---------------- 阶段2：投放后至起爆 ----------------
    t2 = np.linspace(t_t, t_q, 100)
    x2 = x1[-1] + v_u * (t2 - t_t) * np.cos(theta_rad)
    y2 = y1[-1] + v_u * (t2 - t_t) * np.sin(theta_rad)
    z2 = z1[-1] - 0.5 * g * (t2 - t_t)**2   # 自由落体

    # ---------------- 阶段3：起爆后 ----------------
    t3 = np.linspace(t_q, total_duration, 200)
    x3 = np.full_like(t3, x2[-1])
    y3 = np.full_like(t3, y2[-1])
    z3 = z2[-1] - v_sink * (t3 - t_q)

    # ---------------- 触地判定 ----------------
    t_ground_actual = None
    if z3[-1] <= ground_z:
        t_ground_actual = t_q + (z2[-1] - ground_z) / v_sink
        x_ground, y_ground = x2[-1], y2[-1]

    # ---------------- 绘图 ----------------
    fig = plt.figure(figsize=(12, 10))
    gs = GridSpec(3, 1, figure=fig, hspace=0.3)

    ax1, ax2, ax3 = fig.add_subplot(gs[0, 0]), fig.add_subplot(gs[1, 0]), fig.add_subplot(gs[2, 0])

    # 阶段1：蓝色实线
    ax1.plot(t1, x1, "b-", linewidth=2.2, label="投放前（随无人机运动）")
    ax2.plot(t1, y1, "b-", linewidth=2.2)
    ax3.plot(t1, z1, "b-", linewidth=2.2)

    # 阶段2：绿色虚线
    ax1.plot(t2, x2, "g--", linewidth=2.2, label="投放后至起爆前")
    ax2.plot(t2, y2, "g--", linewidth=2.2)
    ax3.plot(t2, z2, "g--", linewidth=2.2)

    # 阶段3：红色点划线
    ax1.plot(t3, x3, "r-.", linewidth=2.2, label="起爆后")
    ax2.plot(t3, y3, "r-.", linewidth=2.2)
    ax3.plot(t3, z3, "r-.", linewidth=2.2)

    # 地面高度线
    ax3.axhline(y=ground_z, color="k", linestyle=":", linewidth=1.5, label="地面高度")

    # ---------------- 关键点标注 ----------------
    # 投放点 S_t
    ax1.scatter(t_t, x1[-1], c="purple", s=90, zorder=5)
    ax1.text(t_t + 0.8, x1[-1], r"$S_t$", color="purple", fontsize=12, fontweight="bold")
    ax2.scatter(t_t, y1[-1], c="purple", s=90, zorder=5)
    ax3.scatter(t_t, z1[-1], c="purple", s=90, zorder=5)

    # 起爆点 C_q
    ax1.scatter(t_q, x2[-1], c="orange", s=90, zorder=5)
    ax1.text(t_q + 0.8, x2[-1], r"$C_q$", color="orange", fontsize=12, fontweight="bold")
    ax2.scatter(t_q, y2[-1], c="orange", s=90, zorder=5)
    ax3.scatter(t_q, z2[-1], c="orange", s=90, zorder=5)

    # 触地点
    if t_ground_actual:
        ax1.scatter(t_ground_actual, x_ground, c="brown", s=90, zorder=5)
        ax1.text(t_ground_actual + 0.8, x_ground, "触地", color="brown", fontsize=12, fontweight="bold")
        ax2.scatter(t_ground_actual, y_ground, c="brown", s=90, zorder=5)
        ax3.scatter(t_ground_actual, ground_z, c="brown", s=90, zorder=5)

    # ---------------- 坐标轴设置 ----------------
    ax1.set_ylabel("$x$ 坐标 (m)")
    ax2.set_ylabel("$y$ 坐标 (m)")
    ax3.set_xlabel("时间 (s)")
    ax3.set_ylabel("$z$ 坐标 (m)")

    fig.suptitle("烟幕弹三阶段运动轨迹图", fontsize=18, y=0.95)

    ax1.legend(loc="upper left", frameon=True, framealpha=0.9, edgecolor="gray")
    ax3.legend(loc="upper right", frameon=True, framealpha=0.9, edgecolor="gray")

    for ax in [ax1, ax2, ax3]:
        ax.grid(True)
        ax.set_xlim(0, total_duration)

    # 阶段说明文字
    stage_text = (
        r"阶段1 ($t \leq t_t$)：随无人机同步运动（水平直线）" + "\n" +
        r"阶段2 ($t_t < t < t_q$)：水平匀速 + 竖直自由落体（曲线）" + "\n" +
        r"阶段3 ($t \geq t_q$)：水平固定 + 竖直匀速下沉（竖直直线）"
    )
    plt.figtext(0.5, 0.01, stage_text, ha="center", fontsize=12,
                bbox=dict(facecolor="white", alpha=0.8, boxstyle="round,pad=0.5"))

    plt.tight_layout()
    plt.savefig("smoke_bomb_trajectory.png", dpi=300, bbox_inches="tight")
    plt.show()


if __name__ == "__main__":
    plot_smoke_bomb_trajectory()
